/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2022-2022. All rights reserved.
 * Description: 性能评估器相关源文件
 * Author: sunzhepeng
 * Create: 2022-8-1
 */

#include "subgraph_star.h"
#include "evaluator_log_def.h"

namespace HcclEvaluator {
SubgraphStar::SubgraphStar(AdjacencyMatrix<> &&matrix) : ConnectedSubgraph(std::forward<AdjacencyMatrix<> &&>(matrix))
{
}

SubgraphStar::~SubgraphStar()
{
}

std::set<AlgType> SubgraphStar::GetSupportAlgType()
{
    return {AlgType::RING, AlgType::RECURSIVE_HD, AlgType::PAIR_WISE};
}

HcclResult SubgraphStar::CalcTimeCost(SubgraphCalcInfo &calcInfo, float &cost)
{
    u32 nodeNum = GetNodeSize();
    u64 byteNum = SIZE_TABLE[calcInfo.dtype] * calcInfo.count;
    AlgType algType = calcInfo.algType;
    SubgraphOpType op = calcInfo.subgraphOpType;

    HockneyModelParam param;
    EVALUATOR_CHK_RET(GetHockneyModelParam(param));
    HockneyModel costModel(param);

    if (algType == AlgType::NOT_SPECIFIED) {
        cost = std::numeric_limits<float>::max();
        for (auto iterAlgType : GetSupportAlgType()) {
            float tmpCost = costModel.CalcCost(nodeNum, byteNum, op, iterAlgType);
            EVALUATOR_INFO("Star topo use alg[%u], cost[%f]\n", static_cast<u32>(iterAlgType), tmpCost);
            if (tmpCost >= 0 && tmpCost < cost) {
                cost = tmpCost;
            }
        }
    } else {
        auto supportAlgSet = GetSupportAlgType();
        if (supportAlgSet.find(algType) != supportAlgSet.end()) {
            cost = costModel.CalcCost(nodeNum, byteNum, op, algType);
        } else {
            EVALUATOR_ERROR("Star topo not support alg[%u]\n", static_cast<u32>(algType));
            return HCCL_E_PARA;
        }
    }

    return HCCL_SUCCESS;
}
} // namespace HcclEvaluator